# coding=utf-8

from Crypto.Cipher import AES
from enum import Enum
import json

PASSWORD = "@Songchao4599!!!"
AES_obj = AES.new(PASSWORD)


class EncodeType(Enum):
    NONE_TYPE = 0
    ASE_TYPE = 1


def encode(msg, encode_type):
    print "EncodeType.NONE_TYPE=", EncodeType.NONE_TYPE
    if encode_type == EncodeType.NONE_TYPE:
        return msg
    elif encode_type == EncodeType.ASE_TYPE:
        return encode_ase(msg)


def encode_ase(msg):
    """
    消息加密
    :param msg: 未加密的消息
    :return: 加密后的消息
    """
    msg_len = len(msg)
    need_len = msg_len % 16
    if need_len != 0:
        need_len = 16 - need_len
        msg += " " * need_len
    return AES_obj.encrypt(msg)


def decode(msg, encode_type):
    """
    消息解密
    :param msg: 客户端传递来的消息
    :param encode_type: 需要解密类型
    :return:
    """
    if encode_type == EncodeType.NONE_TYPE:
        return msg
    elif encode_type == EncodeType.ASE_TYPE:
        return decode_aes(msg)


def decode_aes(msg):
    """
        AES加密消息解密
        :param msg: 未解密的消息
        :return: 解密后的消息
    """
    try:
        decode_msg = AES_obj.decrypt(msg)
    except Exception, e:
        return None

    msg_len = len(decode_msg)
    if msg_len == 0:
        return decode_msg
    try:
        json.loads(decode_msg)
    except Exception, e:
        return None
    return decode_msg
